
#include <apr_general.h>
#include <apr_getopt.h>
#include <apr_strings.h>
#include "aprtools/APRAutoMemoryPool.h"
#include "net/BrokerMessage.h"
#include "net/Broker5Client.h"
#include "resolver/RemoteResolver.h"

#include "utils/RunningWorkProcess.h"

#include "acquire_common.h"
#include "acquire_common_private.h"
#include "acquire_compute.h"

/** A static running work process that can contain a running process
    started from the python script. */
static RunningWorkProcess *r_work_process;

/** Add a key to the local key database from simple text.
 * 
 * @param keyname The name of the key to add to the local key
 * database.
 * 
 * @param keytype The type of the key to import.
 * 
 * @param keydata The data to import the key from.
 * 
 * @param keypurpose The purpose of the key (so, which table to add
 * the key to).
 * 
 * @param docommit Whether to commit() (insert) or update()
 * (update). If the information in the database already exists,
 * commit() fails, if it does not update() fails. */
void acquireComputeAddKey(const char *keyname, const char *keytype, const char *keydata, const char *keypurpose, int docommit)
{
   Key::keypurposetype this_keypurpose;
   if (strcmp(keypurpose, "encryption") == 0) this_keypurpose = Key::KEYPURPOSE_ENCRYPTION;
   if (strcmp(keypurpose, "decryption") == 0) this_keypurpose = Key::KEYPURPOSE_DECRYPTION;
   if (strcmp(keypurpose, "signing") == 0) this_keypurpose = Key::KEYPURPOSE_SIGNING;
   if (strcmp(keypurpose, "checking") == 0) this_keypurpose = Key::KEYPURPOSE_CHECKING;
   Key::keyopenerror terr;
   AbstractKey *a_newkey = Key::tryToGetKeyFromText(keytype, NULL, keydata, keyname, &terr);
   Key::showKeyOpenError(keyname, this_keypurpose, terr);
   if ((a_newkey) && (terr == 0))
   {
      Key tnewkey(state.keydb, keyname, this_keypurpose);
      tnewkey.initFromAbstractKey(state.keydb, *a_newkey, NULL, this_keypurpose);
      if (docommit)
      {
         tnewkey.commit();
      } else
      {
         tnewkey.update();
      }
   } else
   {
      errprint("an error occurred getting the key from the text (error %d)", terr);
   }
}

/** Initialise the compute state.
 *
 * @param hostname The host name of the broker server.
 *
 * @param portno The TCP port number of the broker server.
 *
 * @param sendingkeyname The key (actor) to use for sending data.
 *
 * @param my_name The name of this actor.
 *
 * @param password The password required for opening the key chain.
 *
 * @param recipientkeyname The actor/key name of the recipient.
 *
 * @param keydir The directory containing the keys placed by the SSH
 * system.
 * 
 * @return Zero on failure, one on success. */
int acquireComputeInitState(const char *hostname, int portno, const char *sendingkeyname,
   const char *my_name, const char *password, const char *recipientkeyname, const char *keydir)
{
   r_work_process = NULL;
   apr_status_t status;
   status = apr_pool_create(&(state.mp), (state.apr_iface)->getGlobalPool());
   CheckAPRError(status);
   if (state.mp == NULL)
   {
      errprint("memory pool for results not created");
      return 0;
   }
   state.portno = portno;
   
   if ((hostname == NULL) || (strlen(hostname) == 0))
   {
      errprint("cannot connect to %s hostname", (hostname) ? "an empty" : "a NULL");
      return 0;
   }
   state.hostname = strdup(hostname);
   
   if ((my_name == NULL) || (strlen(my_name) == 0))
   {
      errprint("username cannot be %s", (my_name) ? "left blank" : "NULL"); 
      return 0;
   }
   state.my_name = strdup(my_name);

   if ((sendingkeyname == NULL) || (strlen(sendingkeyname) == 0))
   {
      errprint("sending key name cannot be %s", (sendingkeyname) ? "left blank" : "NULL"); 
      return 0;
   }
   state.operational_senderkeyname = strdup(sendingkeyname);
   state.operational_recipientkeyname = strdup(sendingkeyname);
   
   if ((password == NULL) || (strlen(password) == 0))
   {
      errprint("password cannot be %s", (password) ? "left blank" : "NULL");
      return 0;
   }
   state.password = strdup(password);

   state.keydb = new KeyDB("compute_keydb", NULL, my_name, NULL, 0);
   char *fname = NULL;
   APRAutoMemoryPool pool((state.apr_iface)->getGlobalPool());
   fname = StringUtils::mpprintf(pool.mp(), "%s/%s.enc", (keydir) ? keydir : ".", sendingkeyname);
   importKeyToDatabase(state.keydb, "SymmAES", sendingkeyname, fname, pool.mp(), 1);
   fname = StringUtils::mpprintf(pool.mp(), "%s/%s.dec", (keydir) ? keydir : ".", sendingkeyname);
   importKeyToDatabase(state.keydb, "SymmAES", sendingkeyname, fname, pool.mp(), 1);
   fname = StringUtils::mpprintf(pool.mp(), "%s/%s.sgn", (keydir) ? keydir : ".", sendingkeyname);
   importKeyToDatabase(state.keydb, "PrivateECDSA", sendingkeyname, fname, pool.mp(), 1);
   fname = StringUtils::mpprintf(pool.mp(), "%s/%s.chk", (keydir) ? keydir : ".", sendingkeyname);
   importKeyToDatabase(state.keydb, "PublicECDSA", sendingkeyname, fname, pool.mp(), 1);

   state.resolver = (Resolver *)
      (new RemoteResolver(*(state.keydb),
         state.hostname, state.portno, state.operational_senderkeyname, state.operational_recipientkeyname, state.mp));
   state.registration_state = 1;
   acquireCommonSetInitialised();
   return state.registration_state;
}

/** Add the keys required to decrypt and encrypt work.
 * 
 * @param file_in The file name to use for the input work keys
 * (without extensions).
 * 
 * @param name_in The name to use for the input work keys in the local
 * key database.
 * 
 * @param file_out The file name to use for the output work keys
 * (without extensions).
 * 
 * @param name_out The name to use for the output work keys in the
 * local key database.
 * 
 * @param docommit Whether to commit() (insert) or update()
 * (update). If the information in the database already exists,
 * commit() fails, if it does not update() fails. */
void acquireComputeUpdateUnencryptReencryptKeys(const char *file_in, const char *name_in, const char *file_out, const char *name_out, int docommit)
{
   char *fname = NULL;
   APRAutoMemoryPool pool((state.apr_iface)->getGlobalPool());
   fname = StringUtils::varprintf("%s.dec", file_in);
   importKeyToDatabase(state.keydb, "PrivateRSA", name_in, fname, pool.mp(), docommit);
   free(fname);
   fname = StringUtils::varprintf("%s.enc", file_out);
   importKeyToDatabase(state.keydb, "PublicRSA", name_out, fname, pool.mp(), docommit);
   free(fname);
   fname = StringUtils::varprintf("%s.enc", file_in);
   importKeyToDatabase(state.keydb, "PublicRSA", name_in, fname, pool.mp(), docommit);
   free(fname);
}

/** Open a password file, and decrypt the contents with the given key
 * to yield the plain text password.
 * 
 * @param file_pass The file containing the encrypted password.
 * 
 * @param name_dec The name of the key which should be used to decrypt
 * the password from the file.
 * 
 * @return The plain text password. The string returned is constant
 * because it is held in the results memory pool, which can be cleared
 * using acquireCommonClearResults(). */
const char *acquireComputeGetEncryptedPassword(const char *file_pass, const char *name_dec)
{
   if (file_pass == NULL) return NULL;
   if (name_dec == NULL) return NULL;
   APRAutoMemoryPool pool((state.apr_iface)->getGlobalPool());
   DecryptWrapper *decwrap = (state.keydb)->getDecrypter(name_dec, NULL, 0);
   char *buffer = NULL;
   int bufferlength = 0;
   int retcode = StringUtils::readFileToBuffer(&buffer, &bufferlength, file_pass, pool.mp());
   if (retcode == 0)
   {
      apr_size_t decrypted_length;
      char *decrypted_buffer = decwrap->decryptDataBuffer(buffer, (apr_size_t)bufferlength, &decrypted_length);
      if (buffer) free(buffer);
      char *ret_pntr = (char *)apr_palloc(state.mp, sizeof(char)*(decrypted_length + 1));
      memcpy(ret_pntr, decrypted_buffer, sizeof(char)*decrypted_length);
      if (decrypted_buffer) free(decrypted_buffer);
      ret_pntr[decrypted_length] = '\0';
      return ret_pntr;
   } else return NULL;
}

/** Get the information about the work instance to make a proper local
 * instance of the work.
 * 
 * @return A structure containing the data needed to inform this work
 * instance of the local instance data. */
acquireComputeDataStoreWorkInstanceInfo acquireComputeGetDataStoreForWorkInput()
{
   acquireComputeDataStoreWorkInstanceInfo dswi_info;
   dswi_info.localinstancenumber = -1;
   dswi_info.datastoreid = "";
   if (acquireCommonStateIsIncomplete()) return dswi_info;
   APRAutoMemoryPool pool((state.apr_iface)->getGlobalPool());
   BrokerMessage bmsg(pool.mp());
   bmsg.setRecipient(state.operational_recipientkeyname);
   bmsg.setSender(state.operational_senderkeyname);
   bmsg.addTypeHeaderNOCRLF("REQUEST_WORK", state.my_name);
   if (exchangeMessage(bmsg, pool.mp(), NULL) == 1)
   {
      errprint("message exchange failed");
      return dswi_info;
   }
   // Now we parse the contents of the BrokerMessage in order to extract stuff
   const APRTable *headertable = bmsg.getHeaderTable();
   if ((headertable->get("DS-Name")) == NULL)
   {
      errprint("could not find DS-Name tag in WORK_INSTANCE_INFO (type=%s, processee=%s)", bmsg.getType(), bmsg.getProcessee());
      return dswi_info;
   }
   if ((headertable->get("Local-Instance")) == NULL)
   {
      errprint("could not find work instance number in WORK_INSTANCE_INFO (type=%s, processee=%s)", bmsg.getType(), bmsg.getProcessee());
      return dswi_info;
   }
   dswi_info.localinstancenumber = atoi(headertable->get("Local-Instance"));
   dswi_info.datastoreid = (const char *)apr_pstrdup(state.mp, headertable->get("DS-Name"));
   return dswi_info;
}

/** Start a running work process under this.
 * 
 * @param directory The directory to start the work process in.
 * 
 * @param instanceid The local instance number of the process.
 * 
 * @param instancecount The local instance count, the number of
 * instances in the local group total. */
void acquireComputeStartProcess(const char *directory, int instanceid, int instancecount)
{
   r_work_process = new RunningWorkProcess(directory, instanceid, instancecount, state.mp);
}

/** Watch the running work, and block until the work is cancelled or
 * complete.
 * 
 * @param maxduration The maximum time that this can wait for the work
 * to complete.
 * 
 * @return One if cancelled, zero if completed. */
int acquireComputeWatchForCancelOrComplete(int maxduration)
{
   int cancelled = 0;
   apr_time_t start_time = apr_time_now();
   apr_off_t heartbeats = 0;
   int broker_timeouts_allowed = 3;
   int broker_timeouts_currently = 0;
   while (apr_time_sec(apr_time_now() - start_time) < maxduration)
   {
      if (r_work_process->weAreFinishedYet()) break;
      if ((heartbeats % 5) == 0)
      {
         APRAutoMemoryPool pool((state.apr_iface)->getGlobalPool());
         BrokerMessage bmsg(pool.mp());
         bmsg.setRecipient(state.operational_recipientkeyname);
         bmsg.setSender(state.operational_senderkeyname);
         bmsg.addTypeHeaderNOCRLF("HAPPY_COMPUTE", state.my_name);
         if (exchangeMessage(bmsg, pool.mp(), NULL) == 1)
         {
            errprint("message exchange failed");
            //return 1;
         }
         if (bmsg.isType("DIE_NOW"))
         {
            errprint("process killing because of broker message (work was cancelled?)");
            cancelled = 1;
            break;
         }
         if (bmsg.isType("HAPPY_BROKER"))
         {
            // do nothing, the broker is happy and the work is running
         } else
         {
            // broker is uncommunicative...  if it isn't there,
            // there's no point continuing to run
            broker_timeouts_currently++;
         }
         if (broker_timeouts_allowed == broker_timeouts_currently)
         {
            cancelled = 1;
            errprint("process killing because broker timed out too many times");
            break;
         }
      }
      heartbeats++;
      apr_sleep(APR_USEC_PER_SEC); // 1 second
   }
   errprint("%" APR_OFF_T_FMT " heartbeats", heartbeats);
   if (r_work_process)
      delete r_work_process;
   return cancelled;
}
